from abc import ABC, abstractmethod


class IDocumentExtractNode(ABC):
    @abstractmethod
    def save_context(self, details, workflow_manage=None):
        pass

    @abstractmethod
    def execute(self, document, chat_id, **kwargs):
        pass


class BaseDocumentExtractNode(IDocumentExtractNode):

    def __init__(self, context=None) -> None:
        if not context:
            self.context = {}
        self._context = context.copy()  # Use instance variable instead of setting class attribute

    def __len__(self) -> int:
        return len(self._context)

    def add_to_context(self, key, value):
        """Add new key-value pair to the context."""
        self._context[key] = value

    def get_from_context(self, key):
        """Return value associated with the given key from the context."""
        return self._context.get(key)
    
    def keys_in_context(self):
        return set(self._context.keys())
    

    def clear_all_keys(self):
        """Clears all keys in the context."""
        self._context.clear()

    def save_context(self, details, workflow_manage=None):
        content = details.get('content', '')
        context_key = "content" if content else "details"
        self._add_or_replace(context_key, content or details)
        
        self._answer_text = self._extract_answer(details or '')

    def execute(self, document, chat_id, **kwargs):
        content = kwargs.get('content')
        if content is not None:
            self.add_to_context("content", content)